package cmd

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"time"

	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh"
)

const (
	check_host     = "127.0.0.1"
	check_port     = 2022
	check_user     = "howe"
	check_password = "woainilhj"
)

var checkCmd = &cobra.Command{
	Use:   "check",
	Short: "检查设备状态",
	Long:  `检查设备状态`,
	RunE:  checkDevice,
}

// SSHShellConfig 定义 SSH shell 配置
type SSHShellConfig struct {
	Timeout     time.Duration // 整体操作超时时间
	IdleTimeout time.Duration // 空闲超时时间
}

// SSHShell 封装 SSH shell 操作
type SSHShell struct {
	session *ssh.Session
	stdin   io.WriteCloser
	stdout  io.Reader
	config  SSHShellConfig
}

// NewSSHShell 创建新的 SSH shell
func NewSSHShell(session *ssh.Session, config SSHShellConfig) (*SSHShell, error) {
	stdin, err := session.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("failed to get stdin pipe: %v", err)
	}

	stdout, err := session.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("failed to get stdout pipe: %v", err)
	}

	// 请求伪终端
	if err := session.RequestPty("xterm", 80, 40, ssh.TerminalModes{}); err != nil {
		return nil, fmt.Errorf("failed to request pty: %v", err)
	}

	// 启动 shell
	if err := session.Shell(); err != nil {
		return nil, fmt.Errorf("failed to start shell: %v", err)
	}

	return &SSHShell{
		session: session,
		stdin:   stdin,
		stdout:  stdout,
		config:  config,
	}, nil
}

// ExecuteCommand 执行命令并等待结果
func (s *SSHShell) ExecuteCommand(command string) (string, error) {
	var output bytes.Buffer
	done := make(chan error, 1)
	timeout := make(chan bool, 1)
	idle := make(chan bool, 1)

	// 设置整体超时
	go func() {
		time.Sleep(s.config.Timeout)
		timeout <- true
	}()

	// 读取输出
	go func() {
		buffer := make([]byte, 1024)
		lastActivity := time.Now()

		for {
			n, err := s.stdout.Read(buffer)
			if n > 0 {
				output.Write(buffer[:n])
				lastActivity = time.Now()
			}
			if err != nil {
				if err != io.EOF {
					done <- err
				}
				return
			}

			// 检查空闲超时
			if time.Since(lastActivity) > s.config.IdleTimeout {
				idle <- true
				return
			}
		}
	}()

	// 发送命令
	log.Printf("Executing command: %s", command)
	_, err := fmt.Fprintf(s.stdin, "%s\n", command)
	if err != nil {
		return "", fmt.Errorf("failed to send command: %v", err)
	}

	// 等待结果
	select {
	case err := <-done:
		if err != nil {
			return "", fmt.Errorf("command execution error: %v", err)
		}
	case <-timeout:
		return "", fmt.Errorf("command execution timed out after %v", s.config.Timeout)
	case <-idle:
		return "", fmt.Errorf("command execution idle timeout after %v", s.config.IdleTimeout)
	}

	return strings.TrimSpace(output.String()), nil
}

// Close 关闭 SSH shell
func (s *SSHShell) Close() error {
	return s.session.Close()
}

func checkDevice(cmd *cobra.Command, args []string) error {
	log.Println("Starting device check...")

	// SSH 配置
	config := &ssh.ClientConfig{
		User: check_user,
		Auth: []ssh.AuthMethod{
			ssh.Password(check_password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         10 * time.Second,
		// 添加这些配置可能有助于解决问题
		Config: ssh.Config{
			Ciphers: []string{
				"aes128-ctr", "aes192-ctr", "aes256-ctr",
				"aes128-gcm@openssh.com",
				"chacha20-poly1305@openssh.com",
			},
			KeyExchanges: []string{
				"curve25519-sha256@libssh.org",
				"ecdh-sha2-nistp256",
				"ecdh-sha2-nistp384",
				"ecdh-sha2-nistp521",
				"diffie-hellman-group14-sha1",
			},
		},
	}

	// 连接设备
	address := fmt.Sprintf("%s:%d", check_host, check_port)
	log.Printf("Attempting SSH connection to %s with user %s", address, check_user)
	client, err := ssh.Dial("tcp", address, config)
	if err != nil {
		// 添加网络连接测试
		log.Printf("Testing TCP connection to %s...", address)
		conn, dialErr := net.DialTimeout("tcp", address, 5*time.Second)
		if dialErr != nil {
			log.Printf("TCP connection test failed: %v", dialErr)
		} else {
			conn.Close()
			log.Printf("TCP connection test succeeded")
		}

		return fmt.Errorf("SSH dial error: %v", err)
	}
	defer client.Close()
	log.Println("SSH connection established successfully")

	// 创建新的会话
	session, err := client.NewSession()
	if err != nil {
		return fmt.Errorf("failed to create session: %v", err)
	}
	defer session.Close()

	// 创建 SSH shell
	shell, err := NewSSHShell(session, SSHShellConfig{
		Timeout:     30 * time.Second, // 30秒整体超时
		IdleTimeout: 5 * time.Second,  // 5秒空闲超时
	})
	if err != nil {
		return fmt.Errorf("failed to create shell: %v", err)
	}
	defer shell.Close()

	// 执行命令示例
	log.Println("Searching for serial devices...")
	output, err := shell.ExecuteCommand("ls -l /dev/tty*")
	if err != nil {
		log.Printf("Failed to execute serial device search command: %v", err)
	} else {
		log.Println("Serial device search completed")
		fmt.Printf("Serial devices found:\n%s\n", output)
	}

	// 执行系统更新
	log.Println("Starting system update...")
	output, err = shell.ExecuteCommand("apt update")
	if err != nil {
		log.Printf("Failed to update system: %v", err)
	} else {
		log.Println("System update completed successfully")
		fmt.Printf("System update result:\n%s\n", output)
	}

	log.Println("Device check completed")
	return nil
}

func init() {
	rootCmd.AddCommand(checkCmd)
}
